# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import hashlib
import math
import typing
import uuid

import tornado.gen

import coin.flow.subscriber as flow_subscriber
import coin.flow.topic_map as topic_map

from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_fill_util import FillManager
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.base.rest_order_gateway_base import PrivateExchangeMessage
from coin.exchange.hitbtc_v2.kr_rest.product import HitbtcProduct, HitbtcCurrency
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig

from coin.exchange.hitbtc_v2.ws.ws_private_client import (
    HitbtcWsPrivateClient,
    HitbtcOrderSide,
    HitbtcOrderType,
    HitbtcOrderStatus,
    HitbtcReportType,
)
from coin.exchange.base.rest_order_gateway_base import (
    BaseBalanceOrderGatewayInfo,
    RestOrderGateway,
    exchange_message_handler,
    create_order_internal,
    Order,
    OrderSide,
    OrderType,
    OrderEvent,
    ProductFillElement,
    to_proto_order_side,
)
from coin.proto.coin_query_pb2 import (
    CurrencyBalance,
    AccountBalance,
)


def create_order_from_order_info(order_msg, order_id=None):
  """
  {
    "id": "124432011066",
    "clientOrderId": "pl_2330.6e196c115b5b99f9",
    "symbol": "ETHUSD",
    "side": "sell",
    "status": "canceled",
    "type": "limit",
    "timeInForce": "GTC",
    "quantity": "11.0000",
    "price": "164.260",
    "cumQuantity": "0.0000",
    "createdAt": "2019-05-05T06:51:17.614Z",
    "updatedAt": "2019-05-05T06:51:19.450Z",
    "postOnly": true,
    "reportType": "canceled"
  }
  """
  price_orig = float(order_msg['price'])
  qty_orig = float(order_msg['quantity'])
  order_side = getattr(OrderSide, HitbtcOrderSide(order_msg['side']).name)
  order_type = getattr(OrderType, HitbtcOrderType(order_msg['type']).name)
  exchange_order_id = str(order_msg['id'])
  client_order_id = order_msg['clientOrderId']
  product = HitbtcProduct.FromStrNativeProduct(order_msg['symbol'])
  qty_accum_fill = float(order_msg['cumQuantity'])
  status = HitbtcOrderStatus(order_msg['status'])
  report_type = HitbtcReportType(order_msg['reportType'])
  trade_id = order_msg.get('tradeId')
  if trade_id is not None:
    trade_id = '%s_%s' % (product.symbol, trade_id)
  trade_price = float(order_msg.get('tradePrice', 0))
  trade_quantity = float(order_msg.get('tradeQuantity', 0))

  fill = FillManager(
      qty_orig=qty_orig,
      price_orig=price_orig,
      qty_accum_fill=qty_accum_fill,
      price_avg_fill=price_orig,
      qty_last_fill=qty_accum_fill,
      price_last_fill=price_orig,
  )
  tag = {
      'client_order_id': client_order_id,
      'report_type': report_type,
      'trade_id': trade_id,
      'trade_price': trade_price,
      'trade_quantity': trade_quantity,
  }
  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=tag,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=get_timestamp(),
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                ))
  return order


class HitbtcOrderGatewayInfo(BaseBalanceOrderGatewayInfo):
  CurrencyType = HitbtcCurrency
  ProductType = HitbtcProduct

  def update_balance(self, update_msg: list):
    for balance_info in update_msg:
      currency = HitbtcCurrency.FromStrNativeCurrencyNothrow(balance_info['currency'])
      if currency is None:
        continue

      if currency not in self.currencies:
        continue

      available = float(balance_info['available'])
      hold = float(balance_info['reserved'])
      total = available + hold
      currency_balance = CurrencyBalance(currency=currency.currency,
                                         currency_native=currency.native_currency,
                                         total=total,
                                         available=available,
                                         hold=hold)
      self.set_balance(currency, total, currency_balance)

  def balance_as_proto(self):
    return AccountBalance(exchange='Hitbtc',
                          market_type='Spot',
                          each_balance=self.get_notNone_balances_internal())


class HitbtcOrderGateway(RestOrderGateway):
  _exchange = 'Hitbtc'
  _market_type = 'Spot'
  _api_version = None

  OGInfoType = HitbtcOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    ACCOUNT_BALANCE = 0
    ACTIVE_ORDERS = 1
    ORDER_STATUS_REPORT = 2

  def __init__(self,
               currencies: typing.List[HitbtcCurrency],
               products: typing.List[HitbtcProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__(currencies, products, config, og_config, logger)
    self._private_client = None
    self._scheduler = None

    self._ioloop = None
    self._flow_sub = None
    self._publisher = None

    self._og_config = OrderGatewayConfig(
        account_query_period=1,
        orders_snapshot_query_period=1,
        use_fill_checker=True,
    )
    if og_config is not None:
      self._og_config.MergeFrom(og_config)
    self._uuid = hashlib.md5(str(uuid.uuid1()).encode()).hexdigest()[:16]  # for client order id

  def submittable(self):
    return self.is_ready()

  def _to_client_order_id(self, order):
    client_order_id = '%s.%s' % (order.order_id, self._uuid)
    assert len(client_order_id.split('.')) == 2, client_order_id
    return client_order_id

  @staticmethod
  def _parse_client_order_id(client_order_id):
    pair = client_order_id.split('.')
    if len(pair) == 2:
      res = pair
    else:
      res = (client_order_id, 'unknown')
    return res

  def start(self, queue, ioloop, flow_sub=None):
    topic_string = 'hitbtc_ws_private_1'
    topic_id = topic_map.search_topic_id(topic_string)[0][1]

    self._private_client = HitbtcWsPrivateClient(key_file=self._config.key_file,
                                                 ioloop=self._ioloop,
                                                 notification_callback=self._notification_callback,
                                                 close_callback=self.on_close)
    self._ioloop = ioloop
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)
    self._publisher = queue.get_writer(topic_id=topic_id)
    self._ioloop.add_callback(self._install_periodic_query)
    self._private_client.start()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe(topic_string, self.on_private_exchange_msg)

  def _notification_callback(self, data):
    timestamp = get_timestamp()
    if 'method' not in data:
      self._logger.error('Not a notification message! %s', data)
      return

    method_to_msg_type = {
        'activeOrders': self._MsgType.ACTIVE_ORDERS,
        'report': self._MsgType.ORDER_STATUS_REPORT,
    }

    exchange_message = PrivateExchangeMessage(msg_type=method_to_msg_type[data['method']],
                                              data=data,
                                              received_timestamp=timestamp,
                                              request_params=None,
                                              request_timestamp=None)
    self._publisher.write(timestamp, exchange_message)

  def on_close(self):
    if not self.is_ready():
      self._logger.info('On close triggered!, OG is not ready yet!')
      return

    self._og_ready.set_reinit()
    self._logger.info('Hitbtc ws private client is stopped, will restart soon.')

    self._scheduler.stop()
    self._scheduler.clear()

    self._private_client = self._private_client.clone()
    self._ioloop.call_later(10, self._private_client.start)
    self._ioloop.call_later(12, self._install_periodic_query)

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    while not self._private_client.is_ready():
      self._logger.info('Hitbtc WS private client is not ready yet!')
      yield tornado.gen.sleep(1)

    self._logger.info('Install periodic queryers!')

    # query balance
    query_balance_task = wrap_async_task(
        self._private_client.get_trading_balance,
        msg_type=self._MsgType.ACCOUNT_BALANCE,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(period=self._og_config.account_query_period,
                                      job_func=query_balance_task)

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    # query working orders/dead orders and executions
    query_orders_task = wrap_async_task(
        self._private_client.get_active_orders,
        msg_type=self._MsgType.ACTIVE_ORDERS,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(period=self._og_config.orders_snapshot_query_period,
                                      job_func=query_orders_task)
    self.start_scheduler()

  def _async_submit_order(self, order, post_only=False):
    assert order.order_id is not None, order
    assert order.order_type == OrderType.LIMIT, order
    formatter = self._order_formatter[order.product]
    client_order_id = self._to_client_order_id(order)

    # Note: client order id need to be unique with one trading day according
    # to the api doc.
    params = {
        'symbol': order.product.native_symbol,
        'price': formatter.format_price(order.price),
        'quantity': formatter.format_qty(order.qty),
        'order_side': getattr(HitbtcOrderSide, order.order_side.name),
        'order_type': HitbtcOrderType.LIMIT,
        'client_order_id': client_order_id,
        'post_only': post_only,
    }
    fut = self._private_client.place_new_order(**params)
    return fut

  def _async_cancel_order(self, order):
    if order.tag is not None:
      client_order_id = order.tag['client_order_id']
    else:
      client_order_id = self._to_client_order_id(order)
    fut = self._private_client.cancel_order(client_order_id)
    return fut

  def _handle_submit_response(self, order, msg):
    """ {
      "jsonrpc": "2.0",
      "result": {
        "id": "4345947689",
        "clientOrderId": "57d5525562c945448e3cbd559bd068c4",
        "symbol": "ETHBTC",
        "side": "sell",
        "status": "new",
        "type": "limit",
        "timeInForce": "GTC",
        "quantity": "0.001",
        "price": "0.093837",
        "cumQuantity": "0.000",
        "createdAt": "2017-10-20T12:29:43.166Z",
        "updatedAt": "2017-10-20T12:29:43.166Z",
        "reportType": "new"
      },
      "id": 123
    }
    """
    if msg is None:
      self._logger.error('Submit websocket error %s', msg)
      return OrderEvent.ORDER_ERROR

    if 'error' in msg or 'result' not in msg:
      self._logger.error('Submit order error %s', msg)
      return OrderEvent.ORDER_REJECTED

    resp = msg['result']
    exchange_order_id = str(resp['id'])
    client_order_id = resp['clientOrderId']
    if self._to_client_order_id(order) != client_order_id:
      self._logger.error('Weird! client_order_id is not consistent! %s, %s', order, client_order_id)
      return OrderEvent.ORDER_ERROR
    elif not order.accepted:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
      return OrderEvent.ORDER_ACCEPTED
    else:
      self._logger.debug('order response received already')
      return OrderEvent.ORDER_ACCEPTED

  def _handle_cancel_response(self, order, msg):
    """{
    "jsonrpc": "2.0",
    "result": {
      "id": "4345947689",
      "clientOrderId": "57d5525562c945448e3cbd559bd068c4",
      "symbol": "ETHBTC",
      "side": "sell",
      "status": "canceled",
      "type": "limit",
      "timeInForce": "GTC",
      "quantity": "0.001",
      "price": "0.093837",
      "cumQuantity": "0.000",
      "createdAt": "2017-10-20T12:29:43.166Z",
      "updatedAt": "2017-10-20T12:31:26.174Z",
      "reportType": "canceled"
    },
    "id": 123
    }
    """

    if msg is None:
      self._logger.error('Cancel websocket error!')
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    if 'error' in msg or 'result' not in msg:
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    resp = msg['result']
    exchange_order_id = str(resp['id'])
    if order.internal.exchange_order_id == exchange_order_id:
      if resp['status'] == 'canceled':
        if not self._order_manager.is_dead_order(order):
          self._order_manager.set_order(order, cancel_confirmed=True)
        else:
          self._logger.debug('Order is dead already! %s', order.order_id)
        return OrderEvent.CANCEL_CONFIRMED
    else:
      order.cancel_sent = False
      return OrderEvent.CANCEL_ORDER

  @exchange_message_handler(_MsgType.ACCOUNT_BALANCE)
  def _handle_account_balance(self, exchange_message: PrivateExchangeMessage):
    update_msg = exchange_message.data
    if 'error' in update_msg or 'result' not in update_msg:
      self._logger.error('Invalid balance update msg! %s', update_msg)
      return

    self._og_info.update_balance(update_msg['result'])
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.ACTIVE_ORDERS)
  def _handle_active_orders(self, exchange_message: PrivateExchangeMessage):
    update_msg = exchange_message.data
    timestamp = exchange_message.received_timestamp

    if 'error' in update_msg:
      self._logger.error('query result of active orders! %s', update_msg)
      return

    if not ('result' in update_msg or 'params' in update_msg):
      self._logger.error('Invalid query result! %s', update_msg)
      return

    # From subscription or from active query.
    order_msg_list = update_msg.get('result') or update_msg.get('params', [])

    exchange_order_dict = {}
    for order_msg in order_msg_list:
      order_info = create_order_from_order_info(order_msg)
      exchange_order_id = order_info.internal.exchange_order_id
      client_order_id = order_info.tag.get('client_order_id')
      order_id, og_uuid = self._parse_client_order_id(client_order_id)

      if og_uuid != self._uuid:
        # Already added foreign order.
        order = self._order_manager.get(exchange_order_id=exchange_order_id)
        if order is None:
          # Handle foreign order not seen before.
          self._logger.info('Foreign order (%s, %s) received!', order_id, og_uuid)
          order_id = self._order_id_generator.gen_order_id(is_foreign=True)
          order_info.order_id = order_id
          self._order_manager.add_order(order_info)
      else:
        order_info.order_id = order_id

      if order_info.product in self._og_info.products:
        exchange_order_dict[exchange_order_id] = order_info
      self._update_order(order_info, timestamp)

    self._order_manager.copy_from_exchange_open_orders(exchange_order_dict=exchange_order_dict,
                                                       products=self._og_info.products)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.ORDER_STATUS_REPORT)
  def _handle_order_status_order(self, exchange_message: PrivateExchangeMessage):
    update_msg = exchange_message.data
    timestamp = exchange_message.received_timestamp
    method = update_msg.get('method')
    if method != 'report':
      self._logger.error('Invalid order status report msg! %s', update_msg)
      return
    order_info = create_order_from_order_info(update_msg['params'])
    self._update_order(order_info, timestamp)

  def _update_order(self, order_info, timestamp):
    exchange_order_id = order_info.internal.exchange_order_id
    client_order_id = order_info.tag.get('client_order_id')
    order_id, og_uuid = self._parse_client_order_id(client_order_id)
    status = order_info.internal.status

    if og_uuid != self._uuid:
      # Already added foreign order.
      order = self._order_manager.get(exchange_order_id=exchange_order_id)
      if order is None:
        # Handle foreign order not seen before.
        self._logger.info('Foreign order (%s, %s) received!', order_id, og_uuid)
        order_id = self._order_id_generator.gen_order_id(is_foreign=True)
        order_info.order_id = order_id
        self._order_manager.add_order(order_info)
        order = order_info
    else:
      order = (self._order_manager.get(order_id=order_id)
               or self._order_manager.get(exchange_order_id=exchange_order_id))
      if order is None:
        # This branch should not be reached.
        self._logger.error('Cannot found order! This should not happen! Create a new one, %s',
                           order_info)
        order_id = self._order_id_generator.gen_order_id(is_foreign=True)
        order_info.order_id = order_id
        self._order_manager.add_order(order)
        order = order_info

    if not order.accepted:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)

    if status == HitbtcOrderStatus.CANCELED:
      self._order_manager.set_order(order, cancel_confirmed=True)
    elif status == HitbtcOrderStatus.EXPIRED:
      self._logger.info('Order expired! Auto canceled by exchange! %s', order.order_id)
      self._order_manager.set_order(order, cancel_confirmed=True)
    elif status == HitbtcOrderStatus.FILLED:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, fully_filled=True)
      else:
        self._logger.info('Order marked dead already! %s', order.order_id)
    elif status in (HitbtcOrderStatus.NEW, HitbtcOrderStatus.PARTIALLY_FILLED):
      if self._order_manager.is_dead_order(order):
        self._logger.info('Order is already marked as dead, %s', order)
    elif status == HitbtcOrderStatus.SUSPENDED:
      self._logger.error('Order suspended! %s', order.order_id)
    else:
      self._logger.error('Unsupported order status: %s', status)

    report_type = order_info.tag['report_type']
    if report_type == HitbtcReportType.TRADE:
      trade_id = order_info.tag['trade_id']
      if trade_id not in self._processed_fills:
        self._processed_fills.add(trade_id)
        price_last_fill = order_info.tag['trade_price']
        qty_last_fill = order_info.tag['trade_quantity']
        if math.isclose(qty_last_fill, 0.):
          self._logger.error('Filled qty is 0? %s', order_info.tag)
        last_fill = order.internal.fill.update_by_last_fill(price_last_fill, qty_last_fill)
        assert order in self._order_manager, order
        if last_fill is not None:
          fill_type = self._guess_fill_type(order, price_last_fill, timestamp)
          fill_proto = ProductFillElement(price=last_fill.price_last_fill,
                                          qty=last_fill.qty_last_fill,
                                          symbol=order.product.symbol,
                                          order_id=order.order_id,
                                          side=to_proto_order_side(order.order_side),
                                          fill_id=trade_id,
                                          guess_fill_type=fill_type,
                                          exchange_order_id=order.internal.exchange_order_id)
          self.process_fill(fill_proto, timestamp)
